#include "imgui3d_my_world.h"
#include "implot3d.h"
#include "implot3d_extra.h"
#include <algorithm>
#include <cmath>
#include <deque>
#include <limits>

My3DPlot::My3DPlot(std::string _name) : Control(_name), updateData(nullptr) {}

My3DPlot::~My3DPlot() {}

void My3DPlot::setUpdateDataFunc(std::function<void()> _func)
{
    this->updateData = std::move(_func);
}

void My3DPlot::addOrUpdateBox(const std::string &name, const BoxData &box)
{
    boxes[name] = box;
}

void My3DPlot::addOrUpdateArrow(const std::string &name, const ArrowData &arrow)
{
    arrows[name] = arrow;
}

void My3DPlot::Show()
{
    if (updateData)
    {
        updateData();
    }

    if (boxes.empty())
        return;

    ImGui::Begin(this->name.c_str());

    // 静态变量保持缩放状态与设置
    static float fit_range = 1000.0f;
    static ImVec3 view_center = ImVec3(0, 0, 0);
    static bool center_on_observed = true;
    static bool show_trajectory = false;
    static std::deque<ImVec3> trajectory;

    // Control buttons and options (labels in English)
    if (ImGui::Button("Zoom In"))
    {
        fit_range /= 1.1f;
        if (center_on_observed && boxes.count("observed"))
            view_center = boxes.at("observed").pos;
    }
    ImGui::SameLine();
    if (ImGui::Button("Zoom Out"))
    {
        fit_range *= 1.1f;
        if (center_on_observed && boxes.count("observed"))
            view_center = boxes.at("observed").pos;
    }
    ImGui::SameLine();
    if (ImGui::Button("Reset"))
    {
        fit_range = 1000.0f;
        view_center = ImVec3(0, 0, 0);
    }
    ImGui::Checkbox("Center on red ball", &center_on_observed);
    ImGui::Checkbox("Show trajectory", &show_trajectory);

    static int max_trajectory_len = 120;
    static int pending_trajectory_len = max_trajectory_len;

    ImGui::DragInt("Trajectory Length", &pending_trajectory_len, 1, 10, 1000);
    if (ImGui::Button("Update Length"))
    {
        max_trajectory_len = pending_trajectory_len;
        trajectory.clear(); // 清空当前轨迹
    }

    ImGui::Separator();

    if (ImPlot3D::BeginPlot((this->name + std::string("##3DWorld")).c_str(), ImVec2(-1, 0), ImPlotAxisFlags_AutoFit))
    {
        ImPlot3D::SetupAxesLimits(view_center.x - fit_range, view_center.x + fit_range, view_center.y - fit_range,
                                  view_center.y + fit_range, view_center.z - fit_range, view_center.z + fit_range,
                                  ImPlotCond_Always);

        // 将每个box位置绘制为球形点
        for (const auto &[name, box] : boxes)
            ImPlot3D::SetupAxesLimits(view_center.x - fit_range, view_center.x + fit_range, view_center.y - fit_range,
                                      view_center.y + fit_range, view_center.z - fit_range, view_center.z + fit_range,
                                      ImPlotCond_Always);

        // 将每个box位置绘制为球形点
        for (const auto &[name, box] : boxes)
        {
            float x[1], y[1], z[1];
            x[0] = box.pos.x;
            y[0] = box.pos.y;
            z[0] = box.pos.z;

            ImPlot3D::SetNextMarkerStyle(ImPlot3DMarker_Circle, 5.0f, box.color);
            ImPlot3D::PlotScatter(name.c_str(), x, y, z, 1);
        }

        // 记录并绘制轨迹为点云效果（点云而非连线）
        if (show_trajectory && boxes.count("observed"))
        {
            trajectory.push_back(boxes["observed"].pos);
            if (trajectory.size() > max_trajectory_len)
                trajectory.pop_front();

            std::vector<float> tx, ty, tz;
            for (const auto &p : trajectory)
            {
                tx.push_back(p.x);
                ty.push_back(p.y);
                tz.push_back(p.z);
            }
            ImPlot3D::SetNextMarkerStyle(ImPlot3DMarker_Circle, 2.0f, ImVec4(255, 0, 0, 180));
            ImPlot3D::PlotScatter("Trajectory", tx.data(), ty.data(), tz.data(), (int)trajectory.size());
        }

        // 绘制右下角三轴方向图示
        ImVec3 legend_origin = ImVec3(view_center.x + fit_range * 0.6f, view_center.y - fit_range * 0.6f,
                                      view_center.z - fit_range * 0.6f);
        float len = fit_range * 0.3f;
        // X轴
        {
            float xs[2], ys[2], zs[2];
            ImPlot3D::SetNextLineStyle(ImVec4(255, 0, 0, 255)); // X-axis - red
            xs[0] = legend_origin.x;
            xs[1] = legend_origin.x + len;
            ys[0] = legend_origin.y;
            ys[1] = legend_origin.y;
            zs[0] = legend_origin.z;
            zs[1] = legend_origin.z;
            ImPlot3D::PlotLine("##xaxis_legend", xs, ys, zs, 2);
            ImPlot3D::PlotText("x", xs[1] + len * 0.2, ys[1], zs[1]);
        }
        // Y轴
        {
            float xs[2], ys[2], zs[2];
            ImPlot3D::SetNextLineStyle(ImVec4(0, 255, 0, 255)); // Y-axis - green
            xs[0] = legend_origin.x;
            xs[1] = legend_origin.x;
            ys[0] = legend_origin.y;
            ys[1] = legend_origin.y + len;
            zs[0] = legend_origin.z;
            zs[1] = legend_origin.z;
            ImPlot3D::PlotLine("##yaxis_legend", xs, ys, zs, 2);
            ImPlot3D::PlotText("y", xs[1], ys[1] + len * 0.2, zs[1]);
        }
        // Z轴
        {
            float xs[2], ys[2], zs[2];
            ImPlot3D::SetNextLineStyle(ImVec4(0, 128, 255, 255)); // Z-axis - blue
            xs[0] = legend_origin.x;
            xs[1] = legend_origin.x;
            ys[0] = legend_origin.y;
            ys[1] = legend_origin.y;
            zs[0] = legend_origin.z;
            zs[1] = legend_origin.z + len;
            ImPlot3D::PlotLine("##zaxis_legend", xs, ys, zs, 2);
            ImPlot3D::PlotText("z", xs[1], ys[1], zs[1] + len * 0.2);
        }

        ImPlot3D::EndPlot();
    }

    ImGui::End();
}